import akshare as ak
from MyTT import *
import numpy as np
import pandas as pd
from datetime import datetime
import pymssql
from sqlalchemy import create_engine, distinct, or_, and_
from urllib.parse import quote_plus as urlquote
from configparser import ConfigParser
import warnings

warnings.filterwarnings('ignore')

conf = ConfigParser()
conf.read('env.ini')

sqlserver = ('sqlserver_out', 'sqlserver_in')[0]

host = conf.get(sqlserver, 'host')
port = conf.get(sqlserver, 'port')
user = conf.get(sqlserver, 'user')
password = conf.get(sqlserver, 'password')
schema = conf.get(sqlserver, 'schema')
db = conf.get(sqlserver, 'db')


def get_index_info():  # 返回每个指数的成分股
    sql1 = "SELECT STK_CODE, RIGHT(F_INDEX_CODE, 6) FROM ETF_INDEX_INFO"
    r1 = exec_sql(sql1)
    df1 = pd.DataFrame(r1, columns=['code', 'idx'])
    df1 = df1.loc[df1['idx'] != 'unknowed']
    index_list = list(set(df1.idx.tolist()))

    sql2 = "SELECT INDEX_CODE, CONSTITUENT_CODE FROM index_constituent"
    r2 = exec_sql(sql2)
    df2 = pd.DataFrame(r2, columns=['idx', 'code'])
    res = dict()

    for i, index_code in enumerate(index_list):
        df_p = df2.loc[(df2['idx'] == index_code)]
        if any(df_p.code.str.contains('HK')):
            continue
        if len(df_p) > 0:
            res[index_code] = df_p.code.to_list()

    return res


def get_stocks_data(stocks):
    sql_str = '(' + ','.join(f"'{x}'" for x in stocks) + ')'
    sql = f"SELECT code, datetime, [close], [open], low, high FROM STOCK_DAY_DATA where code in {sql_str}"
    r = exec_sql(sql)
    df = pd.DataFrame(r, columns=['code', 'datetime', 'close', 'open', 'low', 'high'])

    return df


def exec_sql(sql):
    conn = pymssql.connect(host=host, port=port, user=user, password=password, database=db)
    cursor = conn.cursor()
    cursor.execute(sql)
    r = cursor.fetchall()
    cursor.close()
    conn.close()
    return r


def set_date():
    st = "20180101"
    et = "20230915"
    st_d = datetime.strptime(st, '%Y%m%d').date()
    et_d = datetime.strptime(et, '%Y%m%d').date()
    res = {'st': st, 'et': et, 'st_d': st_d, 'et_d': et_d}

    return res


def cal_value():
    tb1, tb2 = 'MB_INFO', 'SCHYD_INFO'
    sql1 = f"SELECT DISTINCT STK_CODE FROM {tb1}"
    sql2 = f"SELECT DISTINCT STK_CODE FROM {tb2}"
    r1 = exec_sql(sql1)
    r2 = exec_sql(sql2)
    in_kd = set([tok[0] for tok in r1])
    in_tk = set([tok[0] for tok in r2])

    date_data = set_date()
    st, et = date_data['st_d'], date_data['et_d']
    df = ak.stock_zh_index_daily(symbol="sz399300")
    df = df.loc[(df.date >= st) & (df.date <= et)]
    df.rename(columns={'date': 'datetime'}, inplace=True)
    df['datetime'] = pd.to_datetime(df['datetime'])
    data = pd.DataFrame(df['datetime'].unique(), columns=['datetime'])
    index_info = get_index_info()
    time_list = df['datetime']
    n_days = len(time_list)

    HS1, LS1 = np.array(n_days * [0]), np.array(n_days * [0])
    HS2, LS2 = np.array(n_days * [0]), np.array(n_days * [0])
    HS3, LS3 = np.array(n_days * [0]), np.array(n_days * [0])
    SS, LS = np.array(n_days * [0]), np.array(n_days * [0])
    st_idx, period = 120, 3

    for i, index in enumerate(index_info):
        print(f'i={i}->len={len(index_info)}->index={index}')
        if index in in_kd and index in in_tk:
            continue

        stockArr = index_info[index]
        stockCounts = len(stockArr)
        is_first = True
        all_df = get_stocks_data(stockArr)

        for code in stockArr:
            df_sig = all_df.loc[all_df.code == code]
            df_sig['datetime'] = pd.to_datetime(df_sig['datetime'])
            df_sig.sort_values('datetime', inplace=True, ignore_index=True)
            if len(df_sig) == 0:
                continue
            data_ = pd.merge(data, df_sig, how='left', on='datetime')
            data_.loc[:, ['code']] = data_.loc[:, ['code']].fillna(method='bfill')
            data_.loc[:, ['close', 'open', 'low', 'high']] = data_.loc[:, ['open', 'high', 'low', 'close']].fillna(
                method='ffill')
            data_.loc[:, ['close', 'open', 'low', 'high']] = data_.loc[:, ['open', 'high', 'low', 'close']].fillna(0)
            high = data_.loc[:, 'high'].to_list()
            low = data_.loc[:, 'low'].to_list()
            close = data_.loc[:, 'close'].to_list()
            open = data_.loc[:, 'open'].to_list()
            low = np.array(low)
            high = np.array(high)
            open = np.array(open)
            close = np.array(close)

            kd_res = cal_kuandu(low, high)
            tk_res = cal_sctk(open, close, low, high)
            if is_first:
                HS1, LS1 = kd_res['HS1'], kd_res['LS1']
                HS2, LS2 = kd_res['HS2'], kd_res['LS2']
                HS3, LS3 = kd_res['HS3'], kd_res['LS3']
                SS, LS = tk_res['SS'], tk_res['LS']
                is_first = False
            else:
                HS1 += kd_res['HS1']
                LS1 += kd_res['LS1']
                HS2 += kd_res['HS2']
                LS2 += kd_res['LS2']
                HS3 += kd_res['HS3']
                LS3 += kd_res['LS3']
                SS += tk_res['SS']
                LS += tk_res['LS']

        K1 = divide(HS1, HS1 + LS1) * 100
        K2 = divide(HS2, HS2 + LS2) * 100
        K3 = divide(HS3, HS3 + LS3) * 100
        A1 = EMA(K1, period)[st_idx:]
        B1 = EMA(K2, period)[st_idx:]
        B2 = EMA(K3, period)[st_idx:]

        HBSTOP1 = SS / stockCounts * 100
        HBSTOP2 = LS / stockCounts * 100
        SD1 = EMA(HBSTOP1, period)[st_idx:]
        SD2 = EMA(HBSTOP2, period)[st_idx:]
        dt_list = time_list[st_idx:]

        date_list = list()

        for i, dt in enumerate(dt_list):
            dt = int(dt.strftime('%Y%m%d'))
            date_list.append(dt)

        kuandu = {'STK_CODE': index[:6], 'DATE_T': date_list, 'A1': A1, 'B1': B1,
                  'B2': B2, 'STK_BD': index[-2:], 'POS_STR': index[-2:]}
        tk = {'STK_CODE': index[:6], 'DATE_T': date_list, 'SD1': SD1, 'SD2': SD2,
              'STK_BD': index[-2:], 'POS_STR': index[-2:]}

        # print(kuandu)
        # print(f'ndays={n_days}->af_ndays={n_days - st_idx}')
        # print(f'K1={len(K1)}->K2={len(K2)}->K3={len(K3)}')
        # print(f'date={len(date_list)}->A1={len(A1)}->B1={len(B1)}->B2={len(B2)}')
        kd_df = pd.DataFrame(kuandu)
        tk_df = pd.DataFrame(tk)
        if index not in in_kd:
            insert_db(kd_df, tb1)
        if index not in in_tk:
            insert_db(tk_df, tb2)
        print(f'{index} data insert ok')


def cal_kuandu(low, high):
    w7, w13, w30 = 7, 13, 30

    vH7, vL7 = REF(HHV(high, w7), 1), REF(LLV(low, w7), 1)
    vH13, vL13 = REF(HHV(high, w13), 1), REF(LLV(low, w13), 1)
    vH30, vL30 = REF(HHV(high, w30), 1), REF(LLV(low, w30), 1)

    HS1, LS1 = high > vH7, low < vL7
    HS2, LS2 = high > vH13, low < vL13
    HS3, LS3 = high > vH30, low < vL30

    # HS1, LS1 = HS1[w7:], LS1[w7:]
    # HS2, LS2 = HS2[w13:], LS2[w13:]
    # HS3, LS3 = HS3[w30:], LS3[w30:]
    HS1, LS1 = HS1.astype(int), LS1.astype(int)
    HS2, LS2 = HS2.astype(int), LS2.astype(int)
    HS3, LS3 = HS3.astype(int), LS3.astype(int)
    res = {'HS1': HS1, 'LS1': LS1, 'HS2': HS2, 'LS2': LS2, 'HS3': HS3, 'LS3': LS3}

    return res


def cal_sctk(open, close, low, high):
    N1, K1, N2, K2 = 13, 6.8, 26, 14

    tj = calc_TJ(close, open, low, high)
    HH1, HH2 = HHV(tj, N1), HHV(tj, N2)
    STOP1, STOP2 = HH1 * (1 - K1 / 100), HH2 * (1 - K2 / 100)
    SS, LS = close < STOP1, close < STOP2
    # SS, LS = SS[N1:], LS[N2:]
    SS, LS = SS.astype(int), LS.astype(int)
    res = {'SS': SS, 'LS': LS}

    return res


def calc_TJ(close, open, low, high):
    ret = []
    for i in range(len(close)):
        _close = close[i]
        _open = open[i]
        _high = high[i]
        _low = low[i]

        v = (_low + (_high - _low) / 2) if _close > _open else (_low - (_high - _low) / 2)
        ret.append(v)

    return np.array(ret)


def insert_db(df, table_name):
    yconnect = create_engine(f'mssql+pymssql://{user}:{urlquote(password)}@{host}:{port}/{db}?charset=utf8')
    pd.io.sql.to_sql(df, table_name, yconnect, schema=schema, if_exists='append', index=False)


def divide(lst1, lst2):
    res = []
    for i in range(len(lst1)):
        a = lst1[i]
        b = lst2[i]
        r = 0 if b == 0 else a / b
        res.append(r)

    return np.array(res)


if __name__ == '__main__':
    cal_value()

